from scipy.special import comb
import numpy as np
import matplotlib.pyplot as plt

def bezier_basis(n, i, t):
    """计算贝塞尔基函数 B_i^n(t)"""
    return comb(n, i) * (1 - t)**(n - i) * t**i


class MPCControllerBezier:
    def __init__(self, A, B, C, N, Q, R, u_min, u_max, reg=1e-5):
        self.A = A
        self.B = B
        self.C = C
        self.N = N  # 预测步长
        self.Q = Q
        self.R = R
        self.u_min = u_min
        self.u_max = u_max
        self.reg = reg  # 正则化参数
        self.n_bezier = 4  # 贝塞尔曲线阶数
        self.measurement_history = []
        self.setpoint_history = []

    def bezier_matrix(self, N, n_bezier):
        """计算贝塞尔基函数矩阵"""
        t_vals = np.linspace(0, 1, N)
        B_matrix = np.zeros((N, n_bezier + 1))
        for k in range(N):
            for i in range(n_bezier + 1):
                B_matrix[k, i] = bezier_basis(n_bezier, i, t_vals[k])
        return B_matrix

    def mpc_control(self, x, r):
        x = x.reshape(-1, 1)
        r = r.reshape(-1, 1)
        n_states = self.A.shape[0]
        n_controls = self.B.shape[1]

        # 计算贝塞尔基函数矩阵
        B_matrix = self.bezier_matrix(self.N, self.n_bezier)

        # 定义优化变量：贝塞尔控制点
        P = np.zeros((self.n_bezier + 1, n_controls))

        # 构建优化问题
        H = np.zeros(((self.n_bezier + 1) * n_controls, (self.n_bezier + 1) * n_controls))  # 15x15
        F = np.zeros(((self.n_bezier + 1) * n_controls, 1))  # 15x1

        G = np.zeros((self.N * n_states, (self.n_bezier + 1) * n_controls))
        E = np.zeros((self.N * n_states, n_states))
        L = np.zeros((self.N * n_states, 1))

        for i in range(self.N):
            for j in range(min(i + 1, self.n_bezier + 1)):  # 修正 j 的范围，避免负幂
                G[i * n_states:(i + 1) * n_states, j * n_controls:(j + 1) * n_controls] \
                    = np.linalg.matrix_power(self.A, i - j) @ self.B * B_matrix[j, 0]  # 修正 `B_matrix` 取值

            E[i * n_states:(i + 1) * n_states, :] = np.linalg.matrix_power(self.A, i + 1)
            L[i * n_states:(i + 1) * n_states, :] = np.linalg.matrix_power(self.A, i + 1) @ x - r

        Q_bar = np.kron(np.eye(self.N), self.Q)
        H += G.T @ Q_bar @ G + self.reg * np.eye(H.shape[0])  # 增加正则化项
        F += (G.T @ Q_bar @ L).reshape(-1, 1)

        # 计算最优贝塞尔控制点 P
        P_opt = -np.linalg.pinv(H) @ F  # 使用 `pinv` 避免不可逆矩阵
        P_opt = P_opt.reshape(self.n_bezier + 1, n_controls)

        # 计算平滑控制输入
        u_opt = B_matrix @ P_opt
        u = u_opt[0]  # 取最开始的控制输入

        return u

    def update(self, setpoint, measurement):
        self.measurement_history.append(measurement)
        self.setpoint_history.append(setpoint)
        u = self.mpc_control(np.array(measurement), np.array(setpoint))
        return u

    def plot_error(self, title='none'):
        self.setpoint_history = np.array(self.setpoint_history)
        self.measurement_history = np.array(self.measurement_history)
        subtitle = ['x', 'y', 'z']
        fig = plt.figure(figsize=(14, 4))
        for i in range(len(self.setpoint_history[0])):
            ax = fig.add_subplot(131+i)
            ax.plot([i / 100 for i in range(len(self.setpoint_history[:, i]))], self.setpoint_history[:, i],
                    label='setpoint')
            ax.plot([i / 100 for i in range(len(self.measurement_history[:, i]))], self.measurement_history[:, i],
                    label='measurement')
            ax.set_xlabel('time')
            ax.set_ylabel('value')
            ax.set_title(subtitle[i]+title)
            ax.legend()
        plt.tight_layout()


if __name__ == '__main__':
    A = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
    B = np.array([[0.5, 0.6, 0.2], [1, 0.1, 0.4], [0.5, 0.5, 0.8]])
    C = np.eye(3)
    Q = np.eye(3) * 2
    R = np.eye(3)
    N = 10
    u_min = -1
    u_max = 1
    reg = 1e-5

    outer_controller = MPCControllerBezier(A, B, C, N, Q, R, u_min, u_max, reg)
    inner_controller = MPCControllerBezier(A, B, C, N, Q, R, u_min, u_max, reg)

    setpoint = np.array([[1], [5], [1]])
    measurement = np.array([[0], [0], [0]])
    x = np.array([[0], [0], [0]])
    x_dot = np.array([[0], [0], [0]])

    for _ in range(10):
        u = outer_controller.update(setpoint, x)
        u_v = inner_controller.update(u, x_dot)
        x_dot = A @ x_dot + B @ u_v.reshape(-1, 1)
        x = A @ x + B @ x_dot.reshape(-1, 1)
        y = C @ x
        print("控制输入 u:", u)
        print('状态 x:', x)
    inner_controller.plot_error()
    outer_controller.plot_error()
    plt.show()
